//
// Created by vivi on 17/04/2018.
//

#pragma once
#include <Eigen/Eigen>
#include <cmath>
#include <utility>
#include <unordered_set>
#include <unordered_map>
#include "../precompiled/precompiled.h"

namespace HairEngine {

	/*
	 * A 3d uniform grid in infinity space
	 */
	template <class T>
	class InfinityGrid3D {
	Public:
		using GridIndex = Eigen::Vector3i;
		using GridIndexAndElement = std::pair<GridIndex, T>;

		struct GridIndexHasher {
			size_t operator()(const GridIndex & index) const {
				//From "Optimized Spatial Hashing for Collision Detection of Deformable Objects"
				size_t v = 73856093 * index.x() + 19349663 * index.y() + 83492791 * index.z();
				return v;
			}
		};

		struct GridIndexAndElementHasher {
			size_t operator()(const GridIndexAndElement & indexAndElement) const {
				GridIndexHasher indexHasher;
				std::hash<T> elementHasher;
				return indexHasher(indexAndElement.first) ^ elementHasher(indexAndElement.second);
			}
		};

		using GridMap = std::unordered_multimap<GridIndex, T, GridIndexHasher>;
		using IndexAndElementSet = std::unordered_set<GridIndexAndElement, GridIndexAndElementHasher>;

		explicit InfinityGrid3D(const Eigen::Vector3f & gridSize_, bool checkDuplicate_ = false):
			gridSize(gridSize_),
			gridSizeInverse(gridSize_.cwiseInverse()),
			checkDuplicate(checkDuplicate_) {

			assert(gridSize.x() > 0.0f);
			assert(gridSize.y() > 0.0f);
			assert(gridSize.z() > 0.0f);

		}

		explicit InfinityGrid3D(float d, bool checkDuplicate_ = false):
			InfinityGrid3D(Eigen::Vector3f(d, d, d), checkDuplicate_) {}

		float dx() const { return gridSize.x(); }
		float dy() const { return gridSize.y(); }
		float dz() const { return gridSize.z(); }

		int getGridIndexX(const float x) const {
			return static_cast<int>(x * gridSizeInverse.x());
		}

		int getGridIndexY(const float y) const {
			return static_cast<int>(y * gridSizeInverse.y());
		}

		int getGridIndexZ(const float z) const {
			return static_cast<int>(z * gridSizeInverse.z());
		}

		/*
		 * Get the correct grid index for the given position, the outFraction is assigned the fraction part of the index
		 * if it is not nullptr
		 */
		GridIndex getGridIndex(const Eigen::Vector3f & pos, Eigen::Vector3f *outFraction = nullptr) const {
			Eigen::Vector3f indexInFloat = pos.cwiseProduct(gridSizeInverse);
			GridIndex index = {
					static_cast<int>(std::floor(indexInFloat.x())),
					static_cast<int>(std::floor(indexInFloat.y())),
					static_cast<int>(std::floor(indexInFloat.z()))
			};

			if (outFraction) {
				*outFraction = {
						indexInFloat.x() - index.x(),
						indexInFloat.y() - index.y(),
						indexInFloat.z() - index.z()
				};
			}

			return index;
		}

		Eigen::AlignedBox3f getAlignedBox(const GridIndex & index) const {
			Eigen::Vector3f pMin = {
					gridSize.x() * index.x(),
					gridSize.y() * index.y(),
					gridSize.z() * index.z()
			};
			Eigen::Vector3f pMax = pMin + gridSize;
			return Eigen::AlignedBox3f(pMin, pMax);
		}

		void insert(const GridIndex & index, const T & element) {
			GridIndexAndElement pair = std::make_pair(index, element);
			if (!checkDuplicate) {
				gridMap.insert(pair);
			}
			else if (indexAndElementSet.count(pair) == 0) {
				indexAndElementSet.insert(pair);
				gridMap.insert(pair);
			}
		}

		/*
		 * Clear all the existing insertion
		 */
		void clear() {
			indexAndElementSet.clear();
			gridMap.clear();
		}

		void rehash(const size_t bucketSize) {
			if (checkDuplicate)
				indexAndElementSet.rehash(bucketSize);
			gridMap.rehash(bucketSize);
		}

		const GridMap & getGridMap() const {
			return gridMap;
		}

		GridMap & getGridMap() {
			return gridMap;
		}

	Protected:
		Eigen::Vector3f gridSize;
		Eigen::Vector3f gridSizeInverse; //Inverse for the grid size, for fast computation

		bool checkDuplicate = false;

		IndexAndElementSet indexAndElementSet;
		GridMap gridMap;
	};
}